import pandas as pd
import torch
import numpy as np

def preprocess():
    train_csv = './data/train_set.csv'
    test_csv = './data/test_a.csv'

    def load_data(file_path, train=True):
        df = pd.read_csv(file_path, sep='\t')
        if train:
            label = df['label'].tolist()
            text = df['text'].tolist()
        else:
            label = None
            text = df['text'].tolist()
        int_text = []
        for i in text:
            tokens = [int(word) for word in i.split()]
            int_text.append(tokens)
        return int_text, label
    
    train_text, label = load_data(train_csv)
    test_text, _ = load_data(test_csv, train=False)
    num_classes = np.unique(label).shape[0]

    token_count= {}
    for line in train_text:
        for word in line:
            if word not in token_count:
                token_count[word] = 1
            else:
                token_count[word] += 1
    for i in range(len(token_count)):
        if i not in token_count:
            padding_idx = i
            break
    max_token_id = max(token_count.keys())

    train_data = {
        'train_text': train_text,
        'train_label': label,
        'num_classes': num_classes,
        'padding_idx': padding_idx,
        'vocab_size': max_token_id + 1
    }
    test_data = {
        'test_text': test_text,
        'num_classes': num_classes,
        'padding_idx': padding_idx,
        'vocab_size': max_token_id + 1
    }

    

    torch.save(train_data, './data/train_data.pt')
    torch.save(test_data, './data/test_data.pt')
    print('Data saved')

if __name__ == '__main__':
    preprocess()